import cv2
import numpy as np

def create_block_difference_mask(gt_img, gen_img, block_size=32):
    """
    创建基于32x32块的差异掩码
    """
    # 确保是灰度图像
    if len(gt_img.shape) == 3:
        gt_img = cv2.cvtColor(gt_img, cv2.COLOR_BGR2GRAY)
    if len(gen_img.shape) == 3:
        gen_img = cv2.cvtColor(gen_img, cv2.COLOR_BGR2GRAY)
    
    # 确保图像大小相同
    if gt_img.shape != gen_img.shape:
        gen_img = cv2.resize(gen_img, (gt_img.shape[1], gt_img.shape[0]))
    
    height, width = gt_img.shape
    mask = np.zeros_like(gt_img)
    
    # 遍历每个块
    for i in range(0, height, block_size):
        for j in range(0, width, block_size):
            y_end = min(i + block_size, height)
            x_end = min(j + block_size, width)
            
            block_gt = gt_img[i:y_end, j:x_end]
            block_gen = gen_img[i:y_end, j:x_end]
            
            diff = cv2.absdiff(block_gt, block_gen)
            if np.mean(diff) > 10:
                mask[i:y_end, j:x_end] = 255
                
    return mask

def apply_block_mask(original_img, mask):
    """
    将块状掩码应用到原始图像
    """
    # 确保是灰度图像
    if len(original_img.shape) == 3:
        original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY)
        
    inv_mask = cv2.bitwise_not(mask)
    masked_img = cv2.bitwise_and(original_img, original_img, mask=inv_mask)
    return masked_img

def visualize_block_differences(gt_img, gen_img, block_size=32):
    """
    可视化块差异
    """
    # 确保是灰度图像
    if len(gt_img.shape) == 3:
        gt_img = cv2.cvtColor(gt_img, cv2.COLOR_BGR2GRAY)
    if len(gen_img.shape) == 3:
        gen_img = cv2.cvtColor(gen_img, cv2.COLOR_BGR2GRAY)
    
    height, width = gt_img.shape
    visualization = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 遍历块
    for i in range(0, height, block_size):
        for j in range(0, width, block_size):
            y_end = min(i + block_size, height)
            x_end = min(j + block_size, width)
            
            block_gt = gt_img[i:y_end, j:x_end]
            block_gen = gen_img[i:y_end, j:x_end]
            
            diff = cv2.absdiff(block_gt, block_gen)
            
            if np.mean(diff) > 5:
                visualization[i:y_end, j:x_end] = [0, 0, 255]  # 红色表示差异
            else:
                visualization[i:y_end, j:x_end] = [0, 255, 0]  # 绿色表示相似
    
    return visualization

def process_images_with_blocks(gt_img, gen_img, block_size=32):
    """
    完整的处理流程
    """
    # 创建块状差异掩码
    diff_mask = create_block_difference_mask(gt_img, gen_img, block_size)
    
    # 应用掩码
    masked_result = apply_block_mask(gen_img, diff_mask)
    
    # 创建可视化结果
    visualization = visualize_block_differences(gt_img, gen_img, block_size)
    
    return masked_result, diff_mask, visualization

# 使用示例
if __name__ == "__main__":
    gt_img = cv2.imread("/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/image/GT.png")
    gen_img = cv2.imread("/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/image/GC_2.png")
    masked_result, diff_mask, visualization = process_images_with_blocks(gt_img, gen_img)
    cv2.imwrite("/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/image/masked_result.png", masked_result)
    cv2.imwrite("/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/image/diff_mask.png", diff_mask)
    cv2.imwrite("/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/image/visualization.png", visualization)
    
